import mitmproxy.http
import csv
from mitmproxy import ctx
import os
import atexit

header = ('url', 'host', 'path', 'methon', 'header', 'request_text', 'status', 'response_text')
with open(r'C:\Users\admin\Desktop\test.csv', 'w+', newline='',encoding='utf8') as f:
    writer = csv.writer(f)
    writer.writerow(header)

class CaptureInfoWriteFile:
    def __init__(self):
        set_proxy()

    def request(self, flow: mitmproxy.http.HTTPFlow):
        info = ctx.log.info
        flow_request = flow.request  # 获取请求对象
        if flow_request.host == '' or flow_request.host is None:
            return
        if flow_request.host == 'test-jzadmin.jzdoctor.com':
            try:
                self.url = flow_request.url  # 接口url
                self.host = flow_request.host  # 域名
                self.path = flow_request.path  # 接口地址
                self.method = flow_request.method  # 请求方式
                self.content = flow_request.content #请求文本
                self.timestamp_start = flow_request.timestamp_start  # 请求开始时间戳
                self.timestamp_end = flow_request.timestamp_end  # 请求结束时间戳
                self.header = flow_request.headers  # 请求头
                self.request_text = flow_request.text  # 请求体
                self.newheader = {}
                for key, value in self.header.items():  # 生成字典样式的header
                    self.newheader[key] = value
            except:
                pass

    def response(self, flow: mitmproxy.http.HTTPFlow):
        info = ctx.log.info
        flow_response = flow.response  # 获取响应对象
        response_status_code = flow_response.status_code  # 响应状态码
        response_text = flow_response.text  # 响应体
        response_content = flow_response.content #响应文本
        response_timestamp_start = flow_response.timestamp_start  # 响应开始时间
        response_timestamp_end = flow_response.timestamp_end  # 响应结束时间
        response_header = flow_response.headers  # 响应头

        # info('响应状态码：'+str(response_status_code))
        # info('响应体：'+str(response_text))

        if "html" in response_text or '<' in response_text or "div" in response_text:
            response_text = 'html网页'

        if "GIF" in response_text or "JPG" in response_text:
            response_text = '图片'

        if "初始化地图" in response_text:
            response_text = '其他'

        # 写入csv文件
        with open(r'C:\Users\admin\Desktop\test.csv', 'a+', newline='',encoding='utf8') as f:
            f_csv = csv.writer(f)
            f_csv.writerows([[self.url, self.host, self.path, self.method, self.newheader, self.request_text, response_status_code, response_text]])

        ## 写入数据库
        # sql = f'''
        # insert into capture_packet(url,host,path,methon,request_header,request_text,status_code,response_text)
        # values('{self.url}','{self.host}','{self.path}','{self.method}','{self.newheader}','{self.request_text}','{response_status_code}','{response_text}')
        # '''
        # Mysql('106.75.225.58', 'root', 'yzp123456', 'auto_data').insert(sql)

    def error(self,flow: mitmproxy.http.HTTPFlow):
        pass

addons = [CaptureInfoWriteFile()]
atexit.register(close_proxy)


if __name__ == '__main__':
    os.system("mitmweb -s mitmproxyRequest.py")